# This is a sample Python script.

# Press Shift+F10 to execute it or replace it with your code.
# Press Double Shift to search everywhere for classes, files, tool windows, actions, and settings.
from __future__ import print_function
import argparse
import os
import sys
import random
import torch
import torch.nn as nn
import torch.nn.parallel
import torch.backends.cudnn as cudnn
cudnn.benchmark = True
cudnn.fastest = True
import torch.optim as optim
import torchvision.utils as vutils
from torch.autograd import Variable
#import os
from misc import *
import models.face_fed as net


from myutils.vgg16 import Vgg16
from myutils import utils
import pdb
import torch.nn.functional as F
#from PIL import Image
from torchvision import transforms

import h5py
from os import listdir
from os.path import isfile, join

parser = argparse.ArgumentParser()
parser.add_argument('--dataset', required=False,
  default='pix2pix_class',  help='')
parser.add_argument('--dataroot', required=False,
  default='', help='path to trn dataset')
parser.add_argument('--valDataroot', required=False,
  default='', help='path to val dataset')
parser.add_argument('--mode', type=str, default='B2A', help='B2A: facade, A2B: edges2shoes')
parser.add_argument('--batchSize', type=int, default=4, help='input batch size')
parser.add_argument('--valBatchSize', type=int, default=120, help='input batch size')
parser.add_argument('--originalSize', type=int,
  default=175, help='the height / width of the original input image')
parser.add_argument('--imageSize', type=int,
  default=128, help='the height / width of the cropped input image to network')
parser.add_argument('--inputChannelSize', type=int,
  default=3, help='size of the input channels')
parser.add_argument('--outputChannelSize', type=int,
  default=3, help='size of the output channels')
parser.add_argument('--ngf', type=int, default=64)
parser.add_argument('--ndf', type=int, default=64)
parser.add_argument('--niter', type=int, default=5000, help='number of epochs to train for')
parser.add_argument('--lrD', type=float, default=0.0002, help='learning rate, default=0.0002')
parser.add_argument('--lrG', type=float, default=0.0002, help='learning rate, default=0.0002')
parser.add_argument('--annealStart', type=int, default=0, help='annealing learning rate start to')
parser.add_argument('--annealEvery', type=int, default=400, help='epoch to reaching at learning rate of 0')
parser.add_argument('--lambdaGAN', type=float, default=0.01, help='lambdaGAN')
parser.add_argument('--lambdaIMG', type=float, default=2.0, help='lambdaIMG')
parser.add_argument('--poolSize', type=int, default=50, help='Buffer size for storing previously generated samples from G')
parser.add_argument('--wd', type=float, default=0.0000, help='weight decay in D')
parser.add_argument('--beta1', type=float, default=0.5, help='beta1 for adam')
parser.add_argument('--netG', default='', help="path to netG (to continue training)")
parser.add_argument('--netD', default='', help="path to netD (to continue training)")
parser.add_argument('--workers', type=int, help='number of data loading workers', default=1)
parser.add_argument('--exp', default='sample', help='folder to output images and model checkpoints')
parser.add_argument('--display', type=int, default=5, help='interval for displaying train-logs')
parser.add_argument('--evalIter', type=int, default=500, help='interval for evauating(generating) images from valDataroot')
opt = parser.parse_args()
print(opt)

from scipy import signal
import h5py
from scipy import signal
import random

#loading kernels mat file
k_filename ='./kernel.mat'
kfp = h5py.File(k_filename)
kernels = np.array(kfp['kernels'])
kernels = kernels.transpose([0,2,1])


vgg = Vgg16()
#utils.init_vgg16('./models/')
state_dict_g = torch.load('VGG_FACE.pth')
new_state_dict_g = {}
new_state_dict_g["conv1_1.weight"]= state_dict_g["0.weight"]
new_state_dict_g["conv1_1.bias"]= state_dict_g["0.bias"]
new_state_dict_g["conv1_2.weight"]= state_dict_g["2.weight"]
new_state_dict_g["conv1_2.bias"]= state_dict_g["2.bias"]
new_state_dict_g["conv2_1.weight"]= state_dict_g["5.weight"]
new_state_dict_g["conv2_1.bias"]= state_dict_g["5.bias"]
new_state_dict_g["conv2_2.weight"]= state_dict_g["7.weight"]
new_state_dict_g["conv2_2.bias"]= state_dict_g["7.bias"]
new_state_dict_g["conv3_1.weight"]= state_dict_g["10.weight"]
new_state_dict_g["conv3_1.bias"]= state_dict_g["10.bias"]
new_state_dict_g["conv3_2.weight"]= state_dict_g["12.weight"]
new_state_dict_g["conv3_2.bias"]= state_dict_g["12.bias"]
new_state_dict_g["conv3_3.weight"]= state_dict_g["14.weight"]
new_state_dict_g["conv3_3.bias"]= state_dict_g["14.bias"]
new_state_dict_g["conv4_1.weight"]= state_dict_g["17.weight"]
new_state_dict_g["conv4_1.bias"]= state_dict_g["17.bias"]
new_state_dict_g["conv4_2.weight"]= state_dict_g["19.weight"]
new_state_dict_g["conv4_2.bias"]= state_dict_g["19.bias"]
new_state_dict_g["conv4_3.weight"]= state_dict_g["21.weight"]
new_state_dict_g["conv4_3.bias"]= state_dict_g["21.bias"]
new_state_dict_g["conv5_1.weight"]= state_dict_g["24.weight"]
new_state_dict_g["conv5_1.bias"]= state_dict_g["24.bias"]
new_state_dict_g["conv5_2.weight"]= state_dict_g["26.weight"]
new_state_dict_g["conv5_2.bias"]= state_dict_g["26.bias"]
new_state_dict_g["conv5_3.weight"]= state_dict_g["28.weight"]
new_state_dict_g["conv5_3.bias"]= state_dict_g["28.bias"]
vgg.load_state_dict(new_state_dict_g)

USE_CUDA = torch.cuda.is_available()
device = torch.device("cuda:0" if USE_CUDA else "cpu")

vgg = torch.nn.DataParallel(vgg,device_ids=[0,2])
vgg.to(device)
#vgg.cuda()

create_exp_dir(opt.exp)
opt.manualSeed=random.randint(1,10000)
random.seed(opt.manualSeed)
torch.manual_seed(opt.manualSeed)
torch.cuda.manual_seed_all(opt.manualSeed)
print("Random Seed:",opt.manualSeed)

opt.dataset='pix2pix_val'
print(opt.dataroot)
dataloader=getLoa